import mongoengine as db
import arrow

from app.db import BaseDocument, BaseEmbeddedDocument
from app.db.models.jobs import Job
from app.db.models.tasks import Task
from app.db.models.users import User


class Notification(BaseEmbeddedDocument):
    email = db.BooleanField(default=False)
    sms = db.BooleanField(default=False)
    wechat = db.BooleanField(default=False)
    inBox = db.BooleanField(default=False)


class Detail(BaseEmbeddedDocument):
    securityStatus = db.StringField(default='')
    taskCategory = db.StringField(default='')
    warningCount = db.LongField(min_value=0, default=0)
    taskType = db.StringField(default='')
    taskId = db.ObjectIdField()
    refId = db.ObjectIdField()
    data = db.DictField(default=dict)


class Alert(BaseDocument):
    meta = {
        'collection': 'alerts',
        'strict': False
    }

    securityStatus = db.StringField(default='')
    taskTriggerType = db.StringField(default='')
    duration = db.LongField(min_value=0, default=0)
    isResolved = db.BooleanField(default=False)
    title = db.StringField(default='')
    summary = db.StringField(default='')
    content = db.StringField(default='')
    detail = db.EmbeddedDocumentField(document_type=Detail)
    taskType = db.StringField(default='')
    domain = db.StringField(default='')
    target = db.StringField(default='')
    jobId = db.ReferenceField(Job)
    refId = db.ReferenceField(User)
    taskId = db.ReferenceField(Task)
    taskCategory = db.StringField(default='')
    alertType = db.StringField(default='')
    taskName = db.StringField(default='')
    uid = db.ReferenceField(User)
    createTime = db.DateTimeField()


    def save(self, *args, **kwargs):
        if not self.createTime:
            self.createTime = arrow.utcnow().datetime
        return super(Alert, self).save(*args, **kwargs)
